from manim import *

class HanoiTower(Scene):
    def construct(self):
        # 基础配置
        n_disks = 3  # 可修改盘片数量
        colors = [BLUE, GREEN, YELLOW, RED, PURPLE]  # 盘片颜色序列
        peg_positions = [-4, 0, 4]  # 三根柱子的x坐标
        
        # 创建柱子标签
        pegs = VGroup(*[
            VGroup(
                Line(ORIGIN, 3*UP, color=GREY),
                Text(f"{chr(65+i)}", font_size=24).next_to(ORIGIN, DOWN)
            ).move_to([x, -1.5, 0])
            for i, x in enumerate(peg_positions)
        ])
        
        # 生成渐变盘片
        disks = VGroup(*[
            Rectangle(
                height=0.6, width=1+0.8*i, 
                fill_color=colors[i%len(colors)], fill_opacity=1,
                stroke_color=WHITE
            ).set_z_index(10-i)
            for i in reversed(range(n_disks))
        ])
        
        # 初始堆叠在A柱
        for i, disk in enumerate(disks):
            disk.move_to([peg_positions[0], -1.5 + 0.6*i, 0])
        
        # 添加初始元素
        self.add(pegs, disks)
        self.wait(1)
        
        # 执行递归动画
        self.generate_movement(n_disks, "A", "C", "B", disks, peg_positions)
        
        # 显示数学公式
        formula = MathTex(r"T(n) = 2T(n-1) + 1 \\ T(1) = 1 \\ \Rightarrow T(n)=2^n-1").to_edge(UR)
        self.play(Write(formula))
        self.wait(2)

    def generate_movement(self, n, source, target, aux, disks, peg_positions, indent=0):
        """递归生成移动动画序列"""
        if n == 0: return
        
        # 显示当前递归层级
        level_text = Text(f"Level {n}", font_size=20).to_edge(UL).shift(RIGHT*indent*0.5)
        self.play(FadeIn(level_text), run_time=0.3)
        
        # 处理上层盘片（修正参数名称）
        self.generate_movement(n-1, source, aux, target, disks, peg_positions, indent+1)
        
        # 移动当前盘片
        moving_disk = disks[-n]
        source_idx = ord(source)-65
        target_idx = ord(target)-65
        
        # 修正点坐标维度（使用传入的peg_positions参数）
        target_x = peg_positions[target_idx]  # 现在可以正确访问
        lift_position = [target_x, moving_disk.get_y()+2, 0]
        final_position = [target_x, -1.5 + 0.6*(n-1), 0]
        
        # 创建移动轨迹
        path = ArcBetweenPoints(
            moving_disk.get_center(),
            lift_position,  # 使用明确的三维坐标
            angle=-PI/2
        )
        
        # 执行动画
        self.play(
            moving_disk.animate.move_to(lift_position),  # 使用修正后的坐标
            Create(path),
            run_time=1.5
        )
        self.play(
            moving_disk.animate.move_to(final_position),  # 使用修正后的坐标
            run_time=1.5
        )
        self.remove(path)
        
        # 处理下层盘片（修正参数名称）
        self.generate_movement(n-1, aux, target, source, disks, peg_positions, indent+1)
        
        # 退出递归层级
        self.play(FadeOut(level_text), run_time=0.3)

# 在命令行执行：manim -pqh hanoi_anime.py HanoiTower